# app.py  —— IndexTTS 单条 + 批量语音克隆
import os
import sys
import time
import shutil
import threading
import zipfile
from pathlib import Path

current_dir = Path(__file__).resolve().parent
sys.path.append(str(current_dir))
sys.path.append(str(current_dir / "indextts"))

import gradio as gr
from indextts.infer import IndexTTS
from tools.i18n.i18n import I18nAuto

i18n = I18nAuto(language="zh_CN")
os.makedirs("outputs/tasks", exist_ok=True)
os.makedirs("prompts", exist_ok=True)

# ---------------- 模型初始化 ----------------
tts = IndexTTS(model_dir="checkpoints", cfg_path="checkpoints/config.yaml")

# ---------------- 推理封装 ----------------
def infer(voice_path: str, text: str, output_path: str | None = None):
    if output_path is None:
        output_path = f"outputs/tasks/spk_{int(time.time()*1000)}.wav"
    tts.infer(voice_path, text, output_path)
    return output_path

# ---------------- 单条生成 ----------------
def gen_single(prompt, text):
    out = infer(prompt, text)
    return gr.update(value=out, visible=True)

# ---------------- 批量生成 ----------------
def scan_prompts():
    """返回 prompts/ 下所有音频文件路径"""
    ext = (".wav", ".flac", ".mp3")
    return [str(p) for p in Path("prompts").rglob("*") if p.suffix.lower() in ext]

def gen_batch(text: str, ref_paths: list, progress=gr.Progress()):
    if not text.strip():
        raise gr.Error("请输入目标文本")
    if not ref_paths:
        raise gr.Error("请至少提供一份参考音频")

    stamp = int(time.time()*1000)
    out_dir = Path(f"outputs/tasks/batch_{stamp}")
    out_dir.mkdir(parents=True, exist_ok=True)
    zip_path = out_dir.with_suffix(".zip")

    progress(0, desc="开始批量生成…")
    for idx, ref in enumerate(ref_paths):
        out_wav = out_dir / f"voice{idx+1}.wav"
        infer(ref, text, str(out_wav))
        progress((idx+1)/len(ref_paths), desc=f"已完成 {idx+1}/{len(ref_paths)}")

    # 打包
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
        for w in out_dir.glob("*.wav"):
            z.write(w, w.name)
    return str(zip_path)

# ---------------- Gradio 界面 ----------------
mutex = threading.Lock()

with gr.Blocks(title="IndexTTS 语音克隆") as demo:
    gr.HTML("""
    <h2><center>IndexTTS: 工业级零样本语音克隆</center></h2>
    <p align="center"><a href='https://www.xiamibaba.com'>Xiamibaba 虾米巴巴</a></p>
    """)

    with gr.Tab("单条生成"):
        with gr.Row():
            prompt_audio = gr.Audio(label="参考音频", sources=["upload", "microphone"], type="filepath")
            input_text_single = gr.Textbox(label="目标文本", lines=3)
            gen_btn = gr.Button("生成", variant="primary")
            out_audio = gr.Audio(label="结果", visible=False)
        gen_btn.click(gen_single, [prompt_audio, input_text_single], out_audio)

    with gr.Tab("批量生成"):
        gr.Markdown("### 步骤：上传/选择多个参考音频 → 输入文本 → 一键打包下载")
        with gr.Row():
            with gr.Column():
                upload_refs = gr.File(label="上传多个参考音频", file_count="multiple", file_types=["audio"])
                scan_btn = gr.Button("扫描 prompts 目录", size="sm")
                ref_listbox = gr.Dropdown(label="已扫描列表（多选）", multiselect=True, interactive=True)
            with gr.Column():
                input_text_batch = gr.Textbox(label="目标文本", lines=4, placeholder="请输入要合成的文本")
                batch_btn = gr.Button("批量生成", variant="primary")
                out_zip = gr.File(label="打包下载", visible=False)

        # 扫描按钮
        scan_btn.click(lambda: gr.update(choices=scan_prompts()), outputs=ref_listbox)

        # 合并上传+扫描，并调用生成
        def gen_batch_wrapper(text, uploaded_files, scanned_files):
            ref_paths = [f.name for f in uploaded_files] if uploaded_files else []
            ref_paths.extend(scanned_files or [])
            return gen_batch(text, ref_paths)

        batch_btn.click(
            fn=gen_batch_wrapper,
            inputs=[input_text_batch, upload_refs, ref_listbox],
            outputs=out_zip
        ).then(lambda: gr.update(visible=True), outputs=out_zip)

if __name__ == "__main__":
    # demo.queue(concurrency_count=5).launch(server_name="127.0.0.1", inbrowser=True)
    demo.queue().launch(server_name="127.0.0.1", inbrowser=True, max_threads=5)